# p is generated as 2-dim
import csv
import datetime
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MaxNLocator
from openpyxl import Workbook
from sko.GA import GA
from sko.tools import set_run_mode

# to monitor optimization process
from osc01.optimization.get_cost import get_cost_with_simulation, get_sampled_duration_and_delivery_time
from osc01.static.constants import Constants, Debugs

iter_counter = 1


def fitness_func(start_correction):
    # print("iterating ...", Constants.N_ACTIVITY, start_correction)
    # print(type(start_correction))
    return get_cost_with_simulation(start_correction=start_correction,
                                    seed=seed,
                                    unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
                                    unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
                                    unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
                                    project_deadline=Debugs.project_deadline
                                    )


# %%

filename_summary = 'summary.xlsx'
if os.path.exists(filename_summary):
    os.remove(filename_summary)

wb_summary = Workbook()
# ws = wb.remove_sheet
ws_summary = wb_summary.active

# prob_mut_dict = {1: 1e-05, 4: 6e-05}
# to try different prob_mut. But, prob_mut is set static. The loop is useless.
# seed is NOT used in this program but in subs.
# it just is a STUD
seed = 1
# for unit_component_late_acceptance_penalty_by_day in [500, 1000, 2000]:
#     for unit_project_late_completion_penalty_by_day in [2000, 5000, 10000]:
sensitivity_table_filename = "sensitivity_table.csv"


def visualize():
    y_history = pd.DataFrame(ga.all_history_Y)
    plt.rcParams['font.family'] = 'Arial'
    plt.figure(figsize=(10, 8), dpi=550)
    # plt.subplots_adjust(bottom=0.2)  # 增加底部空间
    # plt.tight_layout()
    fig, ax = plt.subplots(2, 1)
    ax[0].plot(y_history.index, y_history.values, '.', color='blue')
    ax[0].set_xlabel('Generation')
    ax[0].set_ylabel('Penalty value')
    ax[0].set_title("Genetic algorithm iteration\n" + 'pop=' + str(Debugs.size_pop)
                    + ', max_iter=' + str(Debugs.max_iter) + ', best=' + str(round(best_y[0], 2))
                    # + ',prob_mut=' + str(round(prob_mut, 10))
                    )
    # Y_history.min(axis=1).cummin().plot(kind='line')
    ax[1].bar(range(1, Constants.N_ACTIVITY + 1), np.array(best_x), color='blue')
    ax[1].set_ylim(0, 7)
    ax[1].set_xlabel(f'Activities\' no\n\n'
                     f'CLAP={unit_component_late_acceptance_penalty_by_day}'
                     f'  PLCP={unit_project_late_completion_penalty_by_day}')
    ax[1].set_ylabel('Postponement / day')
    ax[1].set_title("Suggested postponements on the baseline schedule")
    # plt.legend()
    # make locator integer
    plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))
    # fig.savefig('line plot.jpg', bbox_inches='tight', dpi=150)
    fig.savefig('Optimization_CLAP' +
                str(unit_component_late_acceptance_penalty_by_day).zfill(4) +
                '_PLCP' + str(unit_project_late_completion_penalty_by_day).zfill(5) + '.png',
                # bbox_inches='tight',
                dpi=150)
    plt.show()


def getPostponementCost(x):
    # global total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty, total_cost
    avg_component_late_acceptance_penalty, avg_component_late_delivery_penalty, project_late_completion_penalty \
        = get_cost_with_simulation(start_correction=x,
                                   seed=seed,
                                   unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
                                   unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
                                   unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
                                   project_deadline=Debugs.project_deadline,
                                   debug=True
                                   )
    total_cost = avg_component_late_acceptance_penalty \
                 - avg_component_late_delivery_penalty \
                 + project_late_completion_penalty
    return avg_component_late_acceptance_penalty, avg_component_late_delivery_penalty, project_late_completion_penalty, total_cost


# When debugging the format, use a minimal parameter to reduce execution time
# Debugs.max_iter = 1
# Debugs.size_pop = 2
#
with open(sensitivity_table_filename, 'a', newline='') as sensitivity_table_csv_file:
    sensitivity_writer = csv.writer(sensitivity_table_csv_file)
    header = [
        'Unit CLAP/day [$]',
        'Unit PLCP/day [$]',
        'Ratio of unit PLCP to unit CLAP',
        'Avg. TAC of B.S. [$]',
        'Avg. CLAP of B.S. [$]',
        'Avg. CLDP of B.S. [$]',
        'Avg. PLCP of B.S. [$]',
        # 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
        'Avg. suggested postponement [day]',
        'Avg. reduced TAC [$]',
        'Avg. reduced CLAP [$]',
        'Avg. reduced CLDP [$]',
        'Avg. reduced PLCP [$]',
        'Reduction rate'
    ]
    sensitivity_writer.writerow(header)
    for unit_component_late_acceptance_penalty_by_day in [500, 1000, 2000]:
        # for unit_component_late_acceptance_penalty_by_day in [500]:
        for unit_project_late_completion_penalty_by_day in [2000, 5000, 10000]:
            # for unit_project_late_completion_penalty_by_day in [2000]:
            Debugs.unit_component_late_acceptance_penalty_by_day = \
                unit_component_late_acceptance_penalty_by_day
            Debugs.unit_project_late_completion_penalty_by_day = \
                unit_project_late_completion_penalty_by_day
            # prob_mut = prob_mut_dict[seed]
            prob_mut = 1e-06
            # if seed == 6:
            #     prob_mut = 5e-04
            filename = 'Receiving plan under different waiting penalty level' \
                       + str(seed) + '-' + str(round(prob_mut, 10)) + '.xlsx'
            if os.path.exists(filename):
                os.remove(filename)

            wb = Workbook()
            # ws = wb.remove_sheet
            ws = wb.active
            # ws = wb.create_sheet('sheet')
            lb_element = 0
            ub_element = 7
            ws.append(['size_pop=', Debugs.size_pop, 'max_iter=', Debugs.max_iter, 'ub=', ub_element])
            title = [' ', 'seed', 'unit_waiting_penalty', 'unit_project_late_completion_penalty', 'total cost']

            for i in range(1, Constants.N_ACTIVITY + 1):
                title.append('C' + str(i))
            ws.append(title)

            # set_run_mode(get_cost, 'vectorization')
            # set_run_mode(get_cost, 'parallel')
            set_run_mode(fitness_func, 'cached')

            ga = GA(func=fitness_func, n_dim=Constants.N_ACTIVITY,
                    size_pop=Debugs.size_pop,
                    max_iter=Debugs.max_iter,
                    prob_mut=prob_mut,  # probability of mutation
                    lb=[lb_element] * Constants.N_ACTIVITY,
                    ub=[ub_element] * Constants.N_ACTIVITY,
                    # precision=[1] * Constants.N_ACTIVITY)
                    precision=1
                    )

            # ga.Chrom = np.random.randint(0, 1, size=(size_pop, Constants.N_ACTIVITY))
            # manually set the initial population, preset [0 0 0.. 0 0] * 18 * 3 as a chromosome
            inital_pop = np.random.randint(0, 7, size=(Debugs.size_pop, Constants.N_ACTIVITY * 3))
            # set one of individual as [0] * Constants.N_ACTIVITY * 3
            inital_pop[0] = [0] * Constants.N_ACTIVITY * 3
            ga.Chrom = inital_pop

            act_delivery_delay, durations = get_sampled_duration_and_delivery_time(seed=seed)
            base_total_waiting_cost = get_cost_with_simulation(start_correction=[0] * Constants.N_ACTIVITY,
                                                               seed=seed,  # NOT used in the invoked function
                                                               unit_waiting_penalty=Debugs.unit_component_late_acceptance_penalty_by_day,
                                                               unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty_by_day,
                                                               unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty_by_day,
                                                               project_deadline=Debugs.project_deadline
                                                               )
            row = ['base cost', seed, Debugs.unit_component_late_acceptance_penalty_by_day,
                   Debugs.unit_project_late_completion_penalty_by_day]
            row.extend([base_total_waiting_cost])
            row.extend(durations)
            ws_summary.append(row)

            ws.append(row)
            row = ['delivery delay', seed, Debugs.unit_component_late_acceptance_penalty_by_day,
                   Debugs.unit_project_late_completion_penalty_by_day]
            row.extend(' ')
            row.extend(act_delivery_delay)
            ws.append(row)

            start_time = datetime.datetime.now()
            # run ga
            best_x, best_y = ga.run()
            # print("ga.run()", iter_counter)
            iter_counter += 1
            # print('best_x:', best_x, '\n', 'best_y:', best_y)
            lapse_time = (datetime.datetime.now() - start_time).total_seconds()
            ws.append(['lapse', lapse_time])
            row = ['min cost', seed, Debugs.unit_component_late_acceptance_penalty_by_day,
                   Debugs.unit_project_late_completion_penalty_by_day]
            print('best_x=', best_x)
            print('best_y=', best_y)
            row.extend(best_y)
            row.extend(best_x)
            ws.append(row)
            ws_summary.append(row)
            # After solving the problem,plug the solution, best_x to plot
            # Keep temporary data for debug
            print("After solving the problem,plug the solution, best_x to plot")
            # get all the costs of best_x
            avg_component_late_acceptance_penalty, \
            avg_component_late_delivery_penalty, \
            avg_project_late_completion_penalty, \
            total_cost = getPostponementCost(best_x)

            details = ['total cost = ', total_cost,
                       'total_waiting_cost=', round(avg_component_late_acceptance_penalty, 2),
                       'total_late_delivery_penalty=', round(avg_component_late_delivery_penalty, 2),
                       'project_late_completion_penalty=', round(avg_project_late_completion_penalty, 2)]
            ws.append(details)
            ws.append([])
            ws_summary.append(details)
            wb.save(filename)
            wb.close()
            print('seed=', seed, 'lapse_time=', lapse_time)

            visualize()

            # get all the costs of base schedule ([0] * 18)
            avg_component_late_acceptance_penalty_on_base_schedule, \
            avg_component_late_delivery_penalty_on_base_schedule, \
            avg_project_late_completion_penalty_on_base_schedule, \
            avg_total_additional_cost_of_base_schedule = getPostponementCost([0] * 18)

            combined_row = [
                unit_component_late_acceptance_penalty_by_day,  # unit CLAP
                unit_project_late_completion_penalty_by_day,  # # PLCP
                (unit_project_late_completion_penalty_by_day /
                 unit_component_late_acceptance_penalty_by_day),  # Ratio of unit CLAP to unit PLCP
                avg_total_additional_cost_of_base_schedule,  # Avg. TAC
                avg_component_late_acceptance_penalty_on_base_schedule,  # CLAP
                avg_component_late_delivery_penalty_on_base_schedule,  # CLDP
                avg_project_late_completion_penalty_on_base_schedule,  # PLCP
                round(sum(best_x) / len(best_x), 2),  # avg. postponement
                best_y[0],  # reduced TAC, best_y is a ndarray
                round(avg_component_late_acceptance_penalty, 2),  # CLAP
                round(avg_component_late_delivery_penalty, 2),  # CLDP
                round(avg_project_late_completion_penalty, 2),  # PLCP
                round((avg_total_additional_cost_of_base_schedule - best_y[0])
                      / avg_total_additional_cost_of_base_schedule
                      , 2)
            ]
            sensitivity_writer.writerow(combined_row)
        # prob_mut *= 0.4

wb_summary.save(filename_summary)
wb_summary.close()
print("End of optimization")
# base schedule produces an average of additional cost on scenarios

# x = [0] * Constants.N_ACTIVITY
# total_waiting_cost, total_late_delivery_penalty, project_late_completion_penalty \
#     = get_cost_with_simulation(x,
#                                seed=None,
#                                unit_waiting_penalty=Debugs.unit_waiting_penalty,
#                                unit_late_delivery_penalty=Debugs.unit_late_delivery_penalty,
#                                unit_project_late_completion_penalty=Debugs.unit_project_late_completion_penalty,
#                                project_deadline=Debugs.project_deadline,
#                                debug=True
#                                )
# total_cost = total_waiting_cost - total_late_delivery_penalty + project_late_completion_penalty
# print("x=", x, "total_cost=", total_cost)
